from PyQt5.QtWidgets import *
import threading
import sys
from PyQt5.QtCore import *
from PyQt5.QtWidgets import QFileDialog, QMessageBox
from PyQt5.QtGui import *
import face_recognition
import cv2
import os
import numpy as np
from concurrent.futures import ThreadPoolExecutor


class MainWindow(QTabWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle('人脸识别系统')
        self.resize(2400, 1400)  # 设置窗口大小为 2400x1400
        self.setWindowIcon(QIcon("UI_images/logo.png"))
        self.up_img_name = ""
        self.input_fname = ""
        self.source = ''
        self.video_capture = cv2.VideoCapture(0)
        self.stopEvent = threading.Event()
        self.stopEvent.clear()
        self.known_names, self.known_encodings = self.initFaces()
        self.initUI()
        self.set_down()
        self.thread_pool = ThreadPoolExecutor(max_workers=4)  # 线程池

    def initFaces(self):
        known_names = []
        known_encodings = []
        db_folder = "images/db_faces"
        face_imgs = os.listdir(db_folder)
        for face_img in face_imgs:
            face_img_path = os.path.join(db_folder, face_img)
            face_name = face_img.split(".")[0]
            load_image = face_recognition.load_image_file(face_img_path)
            image_face_encoding = face_recognition.face_encodings(load_image)[0]
            known_names.append(face_name)
            known_encodings.append(image_face_encoding)
        return known_names, known_encodings

    def initUI(self):
        # 调整字体大小
        font_v = QFont('楷体', 28)  # 增大字体大小
        generally_font = QFont('楷体', 30)  # 增大字体大小

        # 图片检测界面
        img_widget = QWidget()
        img_layout = QVBoxLayout()
        img_f_title = QLabel("上传人脸图像")
        img_f_title.setAlignment(Qt.AlignCenter)
        img_f_title.setFont(QFont('楷体', 36))  # 增大标题字体大小
        self.img_f_img = QLabel()
        self.img_f_img.setPixmap(QPixmap("UI_images/renlian1.jpg"))
        self.img_f_img.setAlignment(Qt.AlignCenter)
        self.face_name = QLineEdit()
        self.face_name.setPlaceholderText("请输入人脸姓名")
        self.face_name.setFont(generally_font)  # 设置输入框字体大小
        img_up_btn = QPushButton("上传图片")
        img_det_btn = QPushButton("开始上传")
        img_up_btn.clicked.connect(self.up_img)
        img_det_btn.clicked.connect(self.up_db_img)

        # 设置按钮样式
        button_style = """
        QPushButton {
            color: white;
            background-color: #4CAF50;
            border: 2px solid #4CAF50;
            border-radius: 10px;
            padding: 20px;
            margin: 10px;
            font-size: 30px;
        }
        QPushButton:hover {
            background-color: #45a049;
        }
        """
        img_up_btn.setStyleSheet(button_style)
        img_det_btn.setStyleSheet(button_style)

        img_layout.addWidget(img_f_title)
        img_layout.addWidget(self.img_f_img)
        img_layout.addWidget(self.face_name)
        img_layout.addWidget(img_up_btn)
        img_layout.addWidget(img_det_btn)
        img_widget.setLayout(img_layout)

        # 视频检测界面
        video_widget = QWidget()
        video_layout = QVBoxLayout()
        self.video_title2 = QLabel("视频识别区")
        self.video_title2.setFont(font_v)
        self.video_title2.setAlignment(Qt.AlignCenter)
        self.DisplayLabel = QLabel()
        self.DisplayLabel.setPixmap(QPixmap(""))
        self.btn_open_rsmtp = QPushButton("检测摄像头")
        self.btn_open_rsmtp.setFont(font_v)
        self.btn_open = QPushButton("开始识别（选择文件）")
        self.btn_open.setFont(font_v)
        self.btn_open_image = QPushButton("检测图片")
        self.btn_open_image.setFont(font_v)
        self.btn_close = QPushButton("结束检测")
        self.btn_close.setFont(font_v)

        # 设置按钮样式
        self.btn_open_rsmtp.setStyleSheet(button_style)
        self.btn_open.setStyleSheet(button_style)
        self.btn_open_image.setStyleSheet(button_style)
        self.btn_close.setStyleSheet(button_style)

        self.btn_open_rsmtp.clicked.connect(self.open_local)
        self.btn_open.clicked.connect(self.open)
        self.btn_open_image.clicked.connect(self.open_image)
        self.btn_close.clicked.connect(self.close)
        video_layout.addWidget(self.video_title2)
        video_layout.addWidget(self.DisplayLabel)
        video_layout.addWidget(self.btn_open_rsmtp)
        video_layout.addWidget(self.btn_open)
        video_layout.addWidget(self.btn_open_image)
        video_layout.addWidget(self.btn_close)
        video_widget.setLayout(video_layout)

        # 添加选项卡（去掉“关于”页面）
        self.addTab(img_widget, "上传人脸")
        self.addTab(video_widget, '视频检测')
        self.setTabIcon(0, QIcon('UI_images/图片.png'))
        self.setTabIcon(1, QIcon('UI_images/直播.png'))

    def up_img(self):
        # 打开文件选择框，支持JPG和PNG格式
        openfile_name = QFileDialog.getOpenFileName(self, '选择文件', '', 'Image files(*.jpg *.jpeg *.png)')
        # 获取上传的文件名称
        img_name = openfile_name[0]
        if img_name == '':
            pass
        else:
            # 上传之后显示并做归一化处理
            src_img = cv2.imread(img_name)
            if src_img is None:
                QMessageBox.warning(self, "错误", "无法读取图片文件，请检查文件格式！")
                return
            src_img_height = src_img.shape[0]
            src_img_width = src_img.shape[1]
            target_img_height = 800  # 调整图片显示高度
            ratio = target_img_height / src_img_height
            target_img_width = int(src_img_width * ratio)
            # 将图片统一处理到高为800的图片，方便在界面上显示
            target_img = cv2.resize(src_img, (target_img_width, target_img_height))
            # 保存为临时文件
            tmp_img_path = "UI_images/tmp/toup.png" if img_name.lower().endswith('.png') else "UI_images/tmp/toup.jpg"
            cv2.imwrite(tmp_img_path, target_img)
            self.img_f_img.setPixmap(QPixmap(tmp_img_path))
            self.up_img_name = tmp_img_path

    def up_db_img(self):
        # 首先判断该图像是否有一个人脸，多个人脸或者没有人脸都不行
        face_name = self.face_name.text()
        if face_name == "":
            QMessageBox.information(self, "不能为空", "请填写人脸姓名")
        else:
            load_image = face_recognition.load_image_file(self.up_img_name)  # 加载图片
            image_face_encoding = face_recognition.face_encodings(load_image)  # 获得128维特征值
            encoding_length = len(image_face_encoding)  # 获取人脸得数量
            if encoding_length == 0:  # 如果没有人脸，提示用户重新上传
                QMessageBox.information(self, "请重新上传", "当前图片没有发现人脸")
            elif encoding_length > 1:  # 如果人脸有多个，也提示用户重新上传
                QMessageBox.information(self, "请重新上传", "当前图片发现多张人脸")
            else:
                face_encoding = image_face_encoding[0]  # 获取解析得到得人脸数量
                img = cv2.imread(self.up_img_name)  # 将上传得图片保存在db目录下
                img_path = face_name + ('.png' if self.up_img_name.lower().endswith('.png') else '.jpg')
                cv2.imwrite("images/db_faces/" + img_path, img)
                # 上传之后重新对字典进行处理
                self.known_names.append(face_name)
                self.known_encodings.append(face_encoding)
                QMessageBox.information(self, "上传成功", "数据已上传！")

    def open_local(self):
        # 选择录像文件进行读取
        mp4_filename = 0
        self.source = mp4_filename
        # 读取摄像头进行实时得显示
        self.video_capture = cv2.VideoCapture(self.source)
        self.thread_pool.submit(self.display_video)  # 使用线程池

    def open(self):
        # 选择录像文件进行读取
        mp4_fileName, fileType = QFileDialog.getOpenFileName(self, 'Choose file', '', '*.mp4')
        if mp4_fileName:
            # 启动录像文件读取得线程并在画面上实时显示
            self.source = mp4_fileName
            self.video_capture = cv2.VideoCapture(self.source)
            self.thread_pool.submit(self.display_video)  # 使用线程池

    def open_image(self):
        # 选择图片文件进行检测
        img_fileName, _ = QFileDialog.getOpenFileName(self, '选择图片', '', 'Image files(*.jpg *.jpeg *.png)')
        if img_fileName:
            # 读取图片
            image = cv2.imread(img_fileName)
            if image is None:
                QMessageBox.warning(self, "错误", "无法读取图片文件，请检查文件格式！")
                return
            # 进行人脸检测
            rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            face_locations = face_recognition.face_locations(rgb_image)
            face_encodings = face_recognition.face_encodings(rgb_image, face_locations)
            face_names = []
            for face_encoding in face_encodings:
                matches = face_recognition.compare_faces(self.known_encodings, face_encoding, tolerance=0.5)
                if True in matches:
                    first_match_index = matches.index(True)
                    name = self.known_names[first_match_index]
                else:
                    name = "unknown"
                face_names.append(name)
            # 在图片上绘制人脸框和名字
            for (top, right, bottom, left), name in zip(face_locations, face_names):
                cv2.rectangle(image, (left, top), (right, bottom), (0, 0, 255), 4)  # 加粗框线
                cv2.rectangle(image, (left, bottom - 70), (right, bottom), (0, 0, 255), cv2.FILLED)
                font = cv2.FONT_HERSHEY_DUPLEX
                cv2.putText(image, name, (left + 12, bottom - 12), font, 2.0, (255, 255, 255), 2)  # 增大字体
            # 显示检测结果
            image_height = image.shape[0]
            image_width = image.shape[1]
            image_scale = 1000 / image_height  # 调整图片显示大小
            image_resize = cv2.resize(image, (int(image_width * image_scale), int(image_height * image_scale)))
            cv2.imwrite("images/tmp.jpg", image_resize)
            self.DisplayLabel.setPixmap(QPixmap("images/tmp.jpg"))

    def close(self):
        # 点击关闭按钮后重新初始化界面
        self.stopEvent.set()
        self.set_down()

    def set_down(self):
        self.video_capture.release()
        cv2.destroyAllWindows()
        self.DisplayLabel.setPixmap(QPixmap("UI_images/ae862.jpg"))

    def display_video(self):
        # 首先把打开按钮关闭
        self.btn_open.setEnabled(False)
        self.btn_close.setEnabled(True)
        process_this_frame = True
        while True:
            ret, frame = self.video_capture.read()  # 读取摄像头
            if not ret:
                break
            # 降低帧率，每两帧处理一次
            if process_this_frame:
                # 缩小图像分辨率以加快处理速度
                small_frame = cv2.resize(frame, (0, 0), fx=0.5, fy=0.5)
                rgb_small_frame = cv2.cvtColor(small_frame, cv2.COLOR_BGR2RGB)  # 将图像转化为rgb颜色通道
                face_locations = face_recognition.face_locations(rgb_small_frame)  # 获得所有人脸位置
                face_encodings = face_recognition.face_encodings(rgb_small_frame, face_locations)  # 获得人脸特征值
                face_names = []  # 存储出现在画面中人脸的名字
                for face_encoding in face_encodings:  # 和数据库人脸进行对比
                    # 如果当前人脸和数据库的人脸的相似度超过0.5，则认为人脸匹配
                    matches = face_recognition.compare_faces(self.known_encodings, face_encoding, tolerance=0.5)
                    if True in matches:
                        first_match_index = matches.index(True)
                        # 返回相似度最高的作为当前人脸的名称
                        name = self.known_names[first_match_index]
                    else:
                        name = "unknown"
                    face_names.append(name)
            process_this_frame = not process_this_frame
            # 将捕捉到的人脸显示出来
            for (top, right, bottom, left), name in zip(face_locations, face_names):
                # 将检测结果映射回原始分辨率
                top *= 2
                right *= 2
                bottom *= 2
                left *= 2
                cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), 4)  # 加粗框线
                # 加上人名标签
                cv2.rectangle(frame, (left, bottom - 70), (right, bottom), (0, 0, 255), cv2.FILLED)
                font = cv2.FONT_HERSHEY_DUPLEX
                cv2.putText(frame, name, (left + 12, bottom - 12), font, 2.0, (255, 255, 255), 2)  # 增大字体
            # 显示检测结果
            frame_height = frame.shape[0]
            frame_width = frame.shape[1]
            frame_scale = 1000 / frame_height  # 调整图片显示大小
            frame_resize = cv2.resize(frame, (int(frame_width * frame_scale), int(frame_height * frame_scale)))
            # 直接显示图像，避免频繁保存文件
            self.DisplayLabel.setPixmap(QPixmap.fromImage(QImage(frame_resize.data, frame_resize.shape[1], frame_resize.shape[0], QImage.Format_RGB888)))
            if cv2.waitKey(25) & self.stopEvent.is_set() == True:
                self.stopEvent.clear()
                self.DisplayLabel.clear()
                self.btn_close.setEnabled(False)
                self.btn_open.setEnabled(True)
                self.set_down()
                break
        self.btn_open.setEnabled(True)
        self.btn_close.setEnabled(False)
        self.set_down()


if __name__ == "__main__":
    app = QApplication(sys.argv)
    mainWindow = MainWindow()
    mainWindow.show()
    sys.exit(app.exec_())